import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  
import output_function as of
import os
import json
from operatorset import compute_stream_xperiod,apply_bnd_omega_yfreeslip_xperiod
'''守恒格式测试程序。修改好了边界条件'''

'''输入速度，计算涡度。采用中心差分方法。
函数本义是通用的，但此程序中由于没有更新u,v在中间区域的值，所以就在设初值时用了一次
注意，左侧指标是j，右侧指标是i
最上方：[-1, 1:-1]
'''
def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega

#边界涡度处理，采用无滑移条件，顶盖速度恒为1，其余三壁速度恒为零.(太耍赖了，用psi和u,v)
def apply_bnd_omega_(psi, u, v, omega):
    #左右边界周期条件
    omega[:,0]=omega[:,-2]
    omega[:,-1]=omega[:,1]

    #顶部底部涡量直接设为零
    omega[0, 1:-1] =0
    omega[-1, 1:-1] =0

#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
   #边界
    u[:,0]=u[:,-2]
    u[:,-1]=u[:,1]
    u[-1, 1:-1] = u[-2, 1:-1]
    u[0, 1:-1] = u[1, 1:-1]

    v[:,0]=v[:,-2]
    v[:,-1]=v[:,1]
    v[-1, 1:-1] = 0 
    v[0, 1:-1] = 0
    return u,v

#对流算符
def convection_operator(myfun,u,v):
    """
    计算对流算子。
    :param myfun: 二维数组，要计算对流算子的函数。
    :param u: 二维数组，x方向的速度。
    :param v: 二维数组，y方向的速度。
    :param dx: x方向的网格大小。
    :param dy: y方向的网格大小。
    :return: 对流算子的结果。
    """
    """x向对流速度u造成的输运,u \partial_x f"""
    uw_x = 0.5*u[1:-1,1:-1]  * (myfun[1:-1, 2:] - myfun[1:-1, :-2])/ dx
    """v \partial_y f"""
    vw_y = 0.5*v[1:-1,1:-1]  * (myfun[2:,1:-1] - myfun[:-2,1:-1])/ dy

    #守恒形式
    # uw_x =  (u[1:-1, 2:]*myfun[1:-1, 2:] -u[1:-1, :-2]*myfun[1:-1, :-2])/ 2 / dx
    # vw_y =  (v[2:, 1:-1]*myfun[2:, 1:-1] - v[:-2, 1:-1]*myfun[:-2, 1:-1])/2 / dy

    return uw_x,vw_y

#扩散算符
def diff_operator(myfun,nu):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x,diff_y



#漂移驱动
def interchange_term_omega(n_density,beta):
    #beta=Ra*Pr
    linear_flag=False
    if linear_flag:
        drt=beta*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
        return drt
    else:
        drt2=np.zeros_like(n_density)
        n_log_density=np.log(n_density)
        #程序中的x方向差分，但确实是沿极向的
        n_log1=n_log_density[1:-1,2:]
        n_log2=n_log_density[1:-1,:-2]
        #drt2=beta*(n_log[2:,1:-1]-n_log[:-2,1:-1])/2/dy
        drt2[1:-1,1:-1]=beta*(n_log1-n_log2)/2/dy    
        return drt2[1:-1,1:-1]

def interchange_term_n(n_density,psi):
    """phi=-psi"""
    c3=Ra*Pr/Omega_freq
    partial_y_n=(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
    partial_y_psi=(psi[1:-1,2:]-psi[1:-1,:-2])/2/dx
    n_interchange_effect=np.zeros_like(n_density)
    for j in range(1,Ny-1):
        for i in range(1,Nx-1):
            n_interchange_effect[j][i]=zeta*n_density[j][i]*partial_y_psi[j-1][i-1]+c3*partial_y_n[j-1][i-1]
    return n_interchange_effect[1:-1,1:-1]

#涡度向前演化
def forward_omega_FTCS(omega,n_density,u,v):
    uw_x,vw_y =convection_operator(omega,u,v)
    diff_x,diff_y=diff_operator(omega,prandtl)  #扩散项用普朗特数
    omega_SOL_effect=SOL_gamma_omega(n_density,psi) 
    omega_drive_term=interchange_term_omega(n_density,Ra_star*Pr)
    return omega[1:-1, 1:-1] + dt * (-uw_x  -vw_y+omega_drive_term + diff_y + diff_x+omega_SOL_effect)

#密度向前演化
def forward_n_FTCS(n_density,u,v):
    uw_x,vw_y =convection_operator(n_density,u,v)
    diff_x,diff_y=diff_operator(n_density,1)    #归一化的密度扩散系数是1
    n_interchange_term=interchange_term_n(n_density,psi)
    n_SOL_effect=SOL_gamma_n(n_density,psi)
    return n_density[1:-1, 1:-1] + dt * (-uw_x -vw_y+n_interchange_term+n_SOL_effect + diff_y + diff_x)

def SOL_gamma_n(n_density,psi):
    """
    相当于公式中的\Gamma_n,体现SOL效应,在n方程中
    phi=-psi
    alpha_sh0=Omega_freq/L_paral
    """
    c1=L_perp**2/L_paral
    n_SOL_effect=np.zeros_like(n_density)
    for j in range(1,Ny-1):
        for i in range(1,Nx-1):
            n_SOL_effect[j][i]=h_array[j]*(-alpha_sh0*(n_density[j][i]-n_0)-c1*n_density[j][i]*psi[j][i])
    return n_SOL_effect[1:-1, 1:-1]

#
def SOL_gamma_omega(n_density,psi):
    """
    相当于公式中的\Gamma_\omega,体现SOL效应,在omega方程中
    phi=-psi
    """
    c2=L_perp**2/L_paral*Omega_freq
    omega_SOL_effect=np.zeros_like(psi)
    for j in range(1,Ny-1):
        for i in range(1,Nx-1):
            omega_SOL_effect[j][i]=h_array[j]*(-c2*psi[j][i])
    return omega_SOL_effect[1:-1, 1:-1]

def apply_bnd_n_density(n_density):
    """设置密度边条件。对于含SOL的问题,顶边需要设置成第二类边条件"""
    #顶边第二类条件
    n_density[-1,:]=n_density[-2,:]
    #底边第一类条件
    n_density[0,:]=n_down

    #左右边界周期条件
    n_density[:,0]=n_density[:,-2]
    n_density[:,-1]=n_density[:,1]


def diag_n_y_sec(n_density,out_iter):
    
    subfolder='n_density_y_section' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    plt.plot(y,n_density[:,Nx//2])
    plt.xlabel('y')  
    plt.ylabel('n_density') 
    plt.title('Ra={},t={}'.format(Ra,t_diag))  
    plt.grid(True)
    # 定义文件名，包含idiag的值  
    filename_base = 'n_density_y_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def save_simulation_data(n, psi, omega, u, v, out_iter, base_path='./data/SOL/'):
    """
    将模拟数据保存到.npz文件。

    参数:
    n, psi, omega, u, v -- 模拟数据的NumPy数组。
    out_iter -- 当前的时间步。
    base_path -- 保存文件的基础路径。
    """
    # 构建保存路径，包含时间步信息
    save_path = f"{base_path}/simulation_data_t_{out_iter}.npz"
    
    # 确保目录存在
    directory = os.path.dirname(save_path)
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    # 保存数据
    np.savez_compressed(save_path, n=n, psi=psi, omega=omega, u=u, v=v,out_iter=out_iter)



"""设置参数"""
# 给定的网格步长
dx = 0.05  # 例如，我们希望每个网格点之间的距离是 0.01
dy = 0.02  # 同上

# 给定的模拟区域长度
Lx = 5.0  # 模拟区域在 x 方向的长度
Ly = 2.0  # 模拟区域在 y 方向的长度
Ly_in=Ly/2  # 内部
Ly_out=Ly-Ly_in #外部，相当于delta

# 计算网格点数量
Nx = int(Lx / dx) + 1  # 加 1 是为了确保网格点覆盖整个模拟区域，包括边界
Ny = int(Ly / dy) + 1  # 同上

# 打印结果
print(f"Number of grid points in x-direction (Nx): {Nx}")
print(f"Number of grid points in y-direction (Ny): {Ny}")
"""设置密度"""
n_up=n_0=1  #刮削层基底密度
n_down=2    #底边密度
Delta_n=n_down-n_0

x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的。
dx = x[-1] - x[-2]  #网格长
y = np.linspace(0, Ly, Ny)
dy = y[-1] - y[-2]

#设置常数
L_perp=50   #请不要用Lx,因为用过了
L_paral=5500
g_hat=2.4*1e-3
zeta=g_hat*L_perp

#玻姆无量纲下的扩散系数和粘性系数
D_n=0.04
mu=0.1

#由此计算其它参数
Ra_star=g_hat*L_perp**3/D_n/mu
Pr=mu/D_n
Omega_freq=L_perp**2/D_n

alpha_sh0=Omega_freq/L_paral

prandtl=Pr #Pr,普朗特数
Ra=Ra_star #Ra*,瑞利数

dt = 0.00002  #时间步长
ntime=20000  #时间步数
ndiag=100 #每隔多少时间输出一次

xx, yy = np.meshgrid(x, y)  
'''meshgrid生成的矩阵由Ny个长Nx的数列构成，矩阵形状为Ny*Nx。这在索引时造成的顺序是xx[j][i]，如果i是x向指标、j是y向指标的话。
这样，指标数是(i,j)的格点，索引为[j][i]，空间绝对位置(xx[j][i],yy [ j ][i])
'''

# 假设你有一些变量需要保存
myconfig = {
    'dx': dx,
    'dy': dy,
    'Lx': Lx,
    'Ly': Ly,
    'Ly_in':Ly_in,
    'Ly_out':Ly_out,
    'Nx': Nx,
    'Ny': Ny,
    'n_up': n_up,
    'n_0': n_0,
    'n_down': n_down,
    'Delta_n': Delta_n,
    'L_perp': L_perp,
    'L_paral': L_paral,
    'g_hat': g_hat,
    'zeta': zeta,
    'D_n': D_n,
    'mu': mu,
    'Ra_star': Ra_star,
    'Pr': Pr,
    'Omega_freq': Omega_freq,
    'alpha_sh0': alpha_sh0,
    'dt': dt,
    'ntime': ntime,
    'ndiag': ndiag
}
base_folder=f'./data/SOL_Ra={Ra}_Lx_{Lx}_Dn={D_n}_mu={mu}/'
config_filename = f"{base_folder}simulation_config.json"
of.save_to_json(myconfig, config_filename)
print(myconfig)


if __name__ == '__main__':

    """初始化，从零开始"""
    """设置初始流函数场"""
    psi=np.zeros_like(yy)
    u = np.zeros_like(xx)
    v = np.zeros_like(yy)

    """设置初始密度场"""
    tmp_alpha=np.sqrt(alpha_sh0)
    tmp_atad=tmp_alpha*np.tanh(tmp_alpha*Ly_out)
    A_2=1/(1+tmp_atad)
    A_1=-tmp_atad*A_2
    B_2=-np.tanh(tmp_alpha*Ly_out)*A_2
    n_init0_y=n_0*y
    for j in range(0,Ny):
        if y[j]<Ly_in:
            n_init0_y[j]=n_0+Delta_n*(1+A_1*y[j])
        else:
            n_init0_y[j]=n_0+Delta_n*(A_2*np.cosh(tmp_alpha*(y[j]-Ly_in))+B_2*np.sinh(tmp_alpha*(y[j]-Ly_in)))
    n_init0_matrix=np.zeros_like(yy)
    for i in range(0,Nx):
            for j in range(0,Ny):
                n_init0_matrix[j][i]=n_init0_y[j]
    """为密度场添加噪声扰动"""
    n_density=np.copy(n_init0_matrix)
    # 设定噪声的标准差（std_dev），根据需要调整
    std_dev = 0.001  # 这是一个较小的噪声幅度
    #noise = np.random.normal(scale=std_dev, size=n_density.shape)
    noise=np.random.uniform(-std_dev, std_dev, size=n_density.shape)
    for i in range(0,Nx):
        for j in range(0,Ny):
            n_density[j][i]=n_density[j][i]*abs(1+noise[j][i])

    """设置heaviside函数数组,用于计算阶跃函数,用于计算SOL中特有的项"""
    h_array=np.zeros_like(y)
    for j in range(0,Ny):
        h_array[j]=np.heaviside((y[j]-Ly_in),0.5)

    """计算涡度场。依然排除首尾格点"""
    omega = np.zeros_like(xx)
    omega[1:-1, 1:-1] = vel_to_omega(u, v, dx, dy)  
    """omega只需要中间(Nx-2)*(Ny-2)小方块的信息"""
    psi, _ = compute_stream_xperiod(omega[1:-1, 1:-1],Nx,Ny, dx, dy)  
    u,v=psi_to_vel(psi,u,v,dx,dy)
    apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。

    # 选择保存初始数据的路径
    initial_data_path = f'{base_folder}initial_conditions.npz'

    # 调用函数保存初始数据
    of.save_initial_data(n_init0_matrix, xx, yy,x,y, save_path=initial_data_path)

    #进入主循环
    end_signal = False  #终止判断条件。可以再加一个判断条件。
    out_iter = 0    
    #while end_signal==False and out_iter<1000:
    while  out_iter<ntime:
        if out_iter % ndiag==0:
            
            figure_path=f'./SOL_plot_Ra_{Ra}_Lx_{Lx}_Dn={D_n}_mu={mu}'
            of.plot_and_save_v_x_section(x,y, v, Ny//2, out_iter,title_base="V Distribution along X-axis",save_folder_base=figure_path)
            #of.diag_n_origin_pcolor_filepath(n_density, u, v, xx, yy, out_iter, Ra, dt,figure_path)
            of.diag_n_pcolor_filepath(n_density, n_init0_matrix, u, v, xx, yy, out_iter, Ra, dt,figure_path)
            """保存n,psi,omega,u,v到.npz"""
            save_simulation_data(n_density, psi, omega, u, v, out_iter,base_folder)
            print(out_iter)

        out_iter+=1
        """时间向前步进。此时边界涡量会向中心扩散"""
        omega[1:-1, 1:-1] = forward_omega_FTCS(omega,n_density,u,v)  
        n_density[1:-1, 1:-1]=forward_n_FTCS(n_density,u,v)
        """泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。"""
        psi, end_signal = compute_stream_xperiod(omega[1:-1, 1:-1],Nx,Ny, dx, dy, psi_init=psi, tol=1e-4)  
        """根据流函数场计算速度场。"""
        u,v=psi_to_vel(psi,u,v,dx,dy)
        """设置边界涡量。"""
        apply_bnd_omega_(psi, u, v, omega)  
        apply_bnd_n_density(n_density)

    